;; Multiarray descriptors

;; by Konrad Hinsen
;; last updated May 27, 2010

;; Copyright (c) Konrad Hinsen, 2009-2010. All rights reserved.  The use
;; and distribution terms for this software are covered by the Eclipse
;; Public License 1.0 (http://opensource.org/licenses/eclipse-1.0.php)
;; which can be found in the file epl-v10.html at the root of this
;; distribution.  By using this software in any fashion, you are
;; agreeing to be bound by the terms of this license.  You must not
;; remove this notice, or any other, from this software.

(use 'nstools.ns)
(ns+ multiarray.descriptor
  (:clone multiarray.utility))

(set! *warn-on-reflection* true)

;;
;; A MultiArrayDescriptor contains the bookkeeping data of a multiarray,
;; i.e. everything but the data array.
;;

(deftype MultiArrayDescriptor
  [#^ints shape
   #^ints stride
   #^int  offset])

(defmethod print-method MultiArrayDescriptor
  [#^MultiArrayDescriptor descr #^java.io.Writer w]
  (.write w "#:MultiArrayDescriptor{:shape ")
  (print-method (vec (. descr shape)) w)
  (.write w ", :stride ")
  (print-method (vec (. descr stride)) w)
  (.write w ", :offset ")
  (print-method (. descr offset) w)
  (.write w "}"))

;;
;; Make a descriptor for a new array from its shape vector
;;

(defn stride-for-simple
  "return the stride vector for a newly allocated array of the given shape"
  [shape]
  (let [rank (count shape)]
    (if (zero? rank)
      ()
      (let [pstride (conj (subvec shape 1) 1)]
	(reverse (reductions * (reverse pstride)))))))

(defn make-descr
  "return a descriptor for a new contiguous-storage array from its
   shape vector"
  [shape]
  (assert (not-any? neg? shape))
  (let [shape-array  (int-array shape)
	stride-array (int-array (stride-for-simple shape))]
    (new MultiArrayDescriptor shape-array stride-array 0)))

;;
;; Dimension validity tests
;;

(defn assert-rank-not-0
  [#^MultiArrayDescriptor descr]
  (when (zero? (alength #^ints (. descr shape)))
    (throw (Exception. "arrays of rank 0 are not collections"))))

(defn assert-compatible-dim
  [dim rank]
  (when (or (neg? dim) (>= dim rank))
    (throw (Exception. (str "array of rank " rank
			    " does not have dimension " dim)))))

;;
;; Functions working on descriptors
;;

(defn rank
  "returns the number of dimensions of descr"
  [#^MultiArrayDescriptor descr]
  (alength #^ints (. descr shape)))

(defn shape
  "returns a vector containing the number of elements along
   each dimension of descr"
  [#^MultiArrayDescriptor descr]
  (vec (. descr shape)))

(defn shape=
  "returns true if the shapes of two descriptors are equal"
  [#^MultiArrayDescriptor descr1 #^MultiArrayDescriptor descr2]
  (let [#^ints sh1 (. descr1 shape)
	#^ints sh2 (. descr2 shape)]
    (and (= (alength sh1) (alength sh2))
	 (areduce sh1 i t true (and t (= (aget sh1 i) (aget sh2 i)))))))

(defn- remove-from-array
  [#^ints a #^Integer index]
  (let [n  (alength a)
	c  (int-array (dec n))]
    (int-loop [i 0 index]
      (aset c i (aget a i)))
    (int-loop [i (inc index) n]
      (aset c (dec i) (aget a i)))
    c))

(defn sample
  "returns the descriptor for a view sampled along dimension dim
   from index start to index end incrementing by step"
  [#^MultiArrayDescriptor descr #^Integer dim
   #^Integer start #^Integer end #^Integer step]
  (assert (not (zero? step)))
  (let [shape  (aclone #^ints (. descr shape))
        stride (aclone #^ints (. descr stride))
        rank   (alength shape)
	_      (assert-compatible-dim dim rank)
        n      (aget shape dim)
	back?  (neg? step)
        start  (clip-index-to-size n start back?)
        end    (clip-index-to-size n end back?)
        off    (* start (aget stride dim))]
    (aset stride dim #^Integer (* step (aget stride dim)))
    (aset shape dim #^Integer (count-range start end step))
    (new MultiArrayDescriptor shape stride (+ off (. descr offset)))))

(defn slice
  "selects a slice at index along dimension dim"
  [#^MultiArrayDescriptor descr #^Integer dim #^Integer index]
  (let [#^ints shape  (. descr shape)
        #^ints stride (. descr stride)
        rank   (alength shape)
	_      (assert-compatible-dim dim rank)
        n      (aget shape dim)
        index  (check-index-for-size n index)
        off    (* index (aget stride dim))]
    (new MultiArrayDescriptor (remove-from-array shape dim)
	                      (remove-from-array stride dim)
			      (+ off (. descr offset)))))

(defn transpose
  "returns the transpose of the array"
  [#^MultiArrayDescriptor descr dim-permutation]
  (assert (= (sort dim-permutation) (range (rank descr))))
  (new MultiArrayDescriptor (int-array (map (partial nth (. descr shape))
					    dim-permutation))
                            (int-array (map (partial nth (. descr stride))
					    dim-permutation))
			    (. descr offset)))

; The following five functions implement reshape. This is a rather
; complex operation because the compatibility of the data layout with
; the requested shape must be verified while the stride array for the
; new shape is constructed. This happens in two steps: first the shape
; of the input array is simplified as much as possible (functions
; combine-dime and reshape-to-simplest), and then the result is
; reshaped to the requested shape (functions take-factors,
; reshape-from-simplest, reshape). Either step can fail, resulting in
; a nil return value instead of the new descriptor.

(defn- combine-dims
  "Simplifies a list of shape-stride pairs by combining as many as possible."
  [shape-strides]
  (if (< (count shape-strides) 2)
    shape-strides
    (let [[[sh1 str1 :as sh-str-1] [sh2 str2] & more] shape-strides]
      (if (= str1 (* sh2 str2))
	(let [combined [(* sh1 sh2) str2]]
	  (recur (cons combined more)))
	(cons sh-str-1 (combine-dims (rest shape-strides)))))))

(defn- reshape-to-simplest
  "Returns a descriptor for the simplest shape (smallest dimension)
   compatible with reshaping from the given descriptor. The input
   descriptor is for a non-empty array."
  [#^MultiArrayDescriptor descr]
  (let [shape-stride   (combine-dims 
			 (filter (fn [[sh st]] (> sh 1))
				 (map vector (. descr shape) (. descr stride))))
	[shape stride] (if (empty? shape-stride)
			 [[1] [1]]
			 [(map first shape-stride) (map second shape-stride)])
	offset         (. descr offset)]
    (new MultiArrayDescriptor shape stride offset)))

(defn- take-factors
  "Returns the first items of ys such that their product is exactly equal
   to n. If this is not possible, return an empty vector."
  [n xs prod ys]
  (cond (> prod n)  []
	(= prod n)  (vec (concat xs (take-while #(= 1 %) ys)))
	(empty? ys) []
	:else       (let [[y & ys] ys]
		      (recur n (conj xs y) (* prod y) ys))))

(defn- reshape-from-simplest
  "returns a descriptor for a reshape view with shape new-shape based on
   the simplified descriptor descr. Returns nil if reshape is not possible.
   The input descriptor is for a non-empty array."
  [#^MultiArrayDescriptor descr new-shape]
  (let [new-stride
	  (loop [shape      (vec (. descr shape))
		 stride     (vec (. descr stride))
		 new-shape  new-shape
		 new-stride []]
	    (let [[n & shape]  shape
		  [s & stride] stride]
	      (let [sh (take-factors n [] 1 new-shape)
		    st (map #(* s %) (stride-for-simple sh))
		    new-stride (vec (concat new-stride st))]
		(cond (empty? sh)    nil
		      (empty? shape) new-stride
		      :else          (recur shape
					    stride
					    (drop (count sh) new-shape)
					    new-stride)))))]
    (if (nil? new-stride)
      nil
      (new MultiArrayDescriptor (int-array new-shape)
	                        (int-array new-stride)
				(. descr offset)))))

(defn reshape
  "returns a view with a different shape"
  [#^MultiArrayDescriptor descr new-shape]
  (let [n     (apply * (. descr shape))
	new-n (apply * new-shape)]
    (cond (not= n new-n)  nil
	  (zero? n)       (make-descr new-shape)
	  :else           (let [descr (reshape-to-simplest descr)]
			    (reshape-from-simplest descr new-shape)))))

(defn count-elements
  "returns the number of elements in the array"
  [#^MultiArrayDescriptor descr]
  (let [#^ints s (. descr shape)
	r (alength s)]
    (areduce s i p 1 (* p (aget s i)))))

(defn contiguous?
  "returns true if all elements of the array occupy a contiguous index
   range in the data array and are arranged in canonical order"
  [#^MultiArrayDescriptor descr]
  (let [shape        (vec (. descr shape))
        descr-stride (vec (. descr stride))
	cont-stride  (stride-for-simple shape)
	check        (map #(> % 1) shape)
	mismatch     (map not= descr-stride cont-stride)]
    (not (some-true? (map (fn [x y] (and x y)) check mismatch)))))

(defn simple?
  "returns true if the array is contiguous and its data array has the
   minimal required size"
  [#^MultiArrayDescriptor descr #^Integer data-array-length]
  (and (contiguous? descr)
       (= (count-elements descr) data-array-length)))


(set! *warn-on-reflection* false)
